"""
@Project    : cosmo-face
@Module     : ghostnet_face.py
@Author     : HuangJiWen[huangjiwen@haier.com]
@Created    : 2020/8/16 21:07
@Desc       : 基于retina的ghost face
"""

import torch
import torch.nn as nn
import torch.nn.functional as F

from models.retina_face.ghostnet import ghost_net
from models.retina_face.head import BboxHead, ClassHead, LandmarkHead
from models.retina_face.net import FPN, SSH


class RetinaGhostNetFace(nn.Module):

    def __init__(self, cfg=None, phase='train'):
        """
        :param cfg:  Network related settings.
        :param phase: train or test.
        """
        super(RetinaGhostNetFace, self).__init__()
        self.phase = phase

        # backbone
        self.body = ghost_net(width_mult=cfg['width_mult'])

        # self.body = _utils.IntermediateLayerGetter(backbone, cfg['return_layers'])
        in_channels_stage2 = int(cfg['in_channel'] * cfg['width_mult'])
        in_channels_list = [
            in_channels_stage2 * 2,
            int(in_channels_stage2 * 5.6),
            in_channels_stage2 * 8,
        ]
        out_channels = cfg['out_channel']
        self.fpn = FPN(in_channels_list, out_channels)
        self.ssh1 = SSH(out_channels, out_channels)
        self.ssh2 = SSH(out_channels, out_channels)
        self.ssh3 = SSH(out_channels, out_channels)

        self.ClassHead = self._make_class_head(fpn_num=3, inchannels=cfg['out_channel'])
        self.BboxHead = self._make_bbox_head(fpn_num=3, inchannels=cfg['out_channel'])
        self.LandmarkHead = self._make_landmark_head(fpn_num=3, inchannels=cfg['out_channel'])

    def _make_class_head(self, fpn_num=3, inchannels=64, anchor_num=2):
        classhead = nn.ModuleList()
        for i in range(fpn_num):
            classhead.append(ClassHead(inchannels, anchor_num))
        return classhead

    def _make_bbox_head(self, fpn_num=3, inchannels=64, anchor_num=2):
        bboxhead = nn.ModuleList()
        for i in range(fpn_num):
            bboxhead.append(BboxHead(inchannels, anchor_num))
        return bboxhead

    def _make_landmark_head(self, fpn_num=3, inchannels=64, anchor_num=2):
        landmarkhead = nn.ModuleList()
        for i in range(fpn_num):
            landmarkhead.append(LandmarkHead(inchannels, anchor_num))
        return landmarkhead

    def forward(self, inputs):
        """
        x                            ->   torch.Size([2, 3, 608, 608])

        out = self.body(inputs)
        out[1]                       ->   torch.Size([2, 20, 76, 76])
        out[2]                       ->   torch.Size([2, 56, 38, 38])
        out[3]                       ->   torch.Size([2, 80, 19, 19])

        fpn = self.fpn(out)               cfg["out_channel"]
        fpn[0]                       ->   torch.Size([2, 40, 76, 76])
        fpn[1]                       ->   torch.Size([2, 40, 38, 38])
        fpn[2]                       ->   torch.Size([2, 40, 19, 19])

        feature1 = self.ssh1(fpn[0])     ->   torch.Size([2, 40, 76, 76])
        feature2 = self.ssh1(fpn[1])     ->   torch.Size([2, 40, 38, 38])
        feature3 = self.ssh1(fpn[2])     ->   torch.Size([2, 40, 19, 19])
        """
        out = self.body(inputs)

        # FPN
        fpn = self.fpn(out)

        # SSH
        feature1 = self.ssh1(fpn[0])
        feature2 = self.ssh2(fpn[1])
        feature3 = self.ssh3(fpn[2])
        features = [feature1, feature2, feature3]

        bbox_regressions = torch.cat([self.BboxHead[i](feature) for i, feature in enumerate(features)], dim=1)
        classifications = torch.cat([self.ClassHead[i](feature) for i, feature in enumerate(features)], dim=1)
        ldm_regressions = torch.cat([self.LandmarkHead[i](feature) for i, feature in enumerate(features)], dim=1)

        if self.phase == 'train':
            output = (bbox_regressions, classifications, ldm_regressions)
        else:
            output = (bbox_regressions, F.softmax(classifications, dim=-1), ldm_regressions)
        return output


if __name__ == '__main__':
    model = ghost_net()
    model.eval()
    print(model)
    input_x = torch.randn(1, 3, 640, 512)
    y = model(input_x)
    print(y)
